# import qstock as qs
# 沪深300ETF、中证500ETF、红利ETF、创业板ETF、上证50ETF、中证1000ETF、纳斯达克ETF、黄金ETF
# pool_list = ['510300', '510500', '510880', '159915', '510050', '512100', '513100', '518880']
# data = qs.get_price(pool_list, start=start_date, freq='d', fqt=2)
import pandas as pd
import numpy
import quantstats

# 不加这句下面的quantstats.reports.html会报错
numpy.product = numpy.prod


def main():
    WINDOW_SIZE = 20
    df = pd.read_csv("dat1.csv", index_col=0)
    df.drop(columns='中证1000ETF', inplace=True)

    ma20 = df.rolling(WINDOW_SIZE).mean()
    ma20.dropna(inplace=True)
    # 去掉8月23日前的数据
    df_130529 = df.drop(index = df.index[[i for i in range(78)]])
    # print(ma20.shape, df_130529.shape)
    # 计算price / ma20
    price_ma20 = df_130529.div(ma20)
    # 在每一行中找最大的值
    indicators = price_ma20.idxmax(axis=1)
    # 初始资金和持有的份额
    money = 100
    money_list = [money]
    share = 0
    prev_indicator = None
    for idx, indicator in enumerate(indicators):
        if idx == len(indicators) - 1:
            break
        
        if price_ma20[indicator].iloc[idx] <= 1:
            if prev_indicator:
                # 如果price / ma20小于等于1，且现在处于买入某ETF状态，那么卖出
                curr_ETF_price = df_130529[prev_indicator].iloc[idx + 1]
                money = curr_ETF_price * share
                share = 0
                prev_indicator = None
                money_list.append(money)
            else:
                # 没有持有ETF，继续空仓
                money_list.append(money)
                continue
        else:
            if prev_indicator == indicator:
                # price / ma20大于1，且买入标的没有变化，继续持有
                curr_ETF_price = df_130529[indicator].iloc[idx + 1]
                money = curr_ETF_price * share
                money_list.append(money)
                continue
            else:
                if prev_indicator:
                    # 买入的标的发生了变化，转换标的
                    prev_ETF_price = df_130529[prev_indicator].iloc[idx + 1]
                    money = prev_ETF_price * share
                curr_ETF_price = df_130529[indicator].iloc[idx + 1]
                share = money / curr_ETF_price
                prev_indicator = indicator
                money_list.append(money)
    print(money, len(money_list), ma20.shape)
    df_130529['MyStrategy'] = [m / 100.0 for m in money_list]
    print(df_130529.head())
    df_130529.index = pd.to_datetime(df_130529.index)
    quantstats.reports.html(df_130529['MyStrategy'], benchmark=df_130529['沪深300ETF'], title="MA20策略报告", output="MA20_report.html")

if __name__ == "__main__":
    main()


    